package ws

import (
    gorilla "github.com/gorilla/websocket"
    "github.com/kataras/iris/v12"
    "github.com/kataras/iris/v12/websocket"
    "github.com/kataras/neffos"
    irisgorilla "github.com/kataras/neffos/gorilla"
    "log"
    "net/http"
)

type IrisWebSocketGoroutine struct {
    ws             *neffos.Server
    ReceiveMessage chan []byte // 接收消息的通道
    SendMessage    chan string // 推送消息的通道
}

// NewWebSocketGoroutine 初始化 WebSocket
func NewWebSocketGoroutine(app *iris.Application) *IrisWebSocketGoroutine {

    i := &IrisWebSocketGoroutine{
        ReceiveMessage: make(chan []byte),
        SendMessage:    make(chan string),
    }

    // 使用gorilla/websocket模块实现WebSocket
    u := irisgorilla.Upgrader(gorilla.Upgrader{
        // 允许跨域
        CheckOrigin: func(r *http.Request) bool {
            return true
        },
    })

    i.ws = websocket.New(u, websocket.Events{
        websocket.OnNativeMessage: func(nsConn *websocket.NSConn, msg websocket.Message) error {
            // 将接收到的消息发送到通道中
            i.ReceiveMessage <- msg.Body
            return nil
        },
    })

    // 监听服务是否链接
    i.ws.OnConnect = func(c *websocket.Conn) error {
        log.Printf("[%s] 服务已连接!", c.ID())
        return nil
    }

    // 监听链接是否断开
    i.ws.OnDisconnect = func(c *websocket.Conn) {
        log.Printf("[%s] 链接已断开!", c.ID())
    }

    i.ws.OnUpgradeError = func(err error) {
        log.Printf("Upgrade Error: %v", err)
    }

    // 开启协程，监听通道是否接收到消息
    go func() {
        for true {
            select {
            case msgChan := <-i.ReceiveMessage:
                log.Printf("接收到客户端消息: %s", msgChan)
            }
        }
    }()

    // 开启协程，监听通道是否要推送消息
    go func() {
        for true {
            select {
            case msgChan := <-i.SendMessage:
                // 推送消息
                i.ws.Broadcast(nil, websocket.Message{
                    Body:     []byte(msgChan), // 消息体
                    IsNative: true,            // 使用自己的格式传递数据
                })
            }
        }
    }()

    // 设置WebSocket链接地址
    app.Get("/websocket", websocket.Handler(i.ws))

    return i
}
